import os
import json
import process_file as pf
import csv
import csv_related as cr

folder_path_root = '../result/json_train_valid_test/'
save_path_root = '../result/csv_train_valid_test/'


subdirectories = pf.get_first_level_subdirectories(folder_path_root)
for dir_path in subdirectories:
    dir_name = pf.get_last_directory_name(dir_path)
    save_path = f'{save_path_root}{dir_name}.csv'
    # 确保输出目录存在
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    with open(save_path, mode='w', newline='', encoding='utf-8') as csvfile:
            # 创建CSV写入器
            csvwriter = csv.writer(csvfile)
            # 写入CSV文件的表头
            csvwriter.writerow(['input_function', 'output_doctest'])  # 替换为你需要的列名


            # 遍历文件夹中的所有文件
            for filename in os.listdir(dir_path):
                file_path = os.path.join(dir_path, filename)
                if os.path.isfile(file_path) and filename.endswith('.json'):
                    list_dict_info = pf.read_json(file_path)

                    for dict_info in list_dict_info:
                        # 提示
                        input_question = f'//Rust Function under test\n' \
                                         f'{dict_info["FUT_full"]}\n' \
                                         f'//Test Function'
                        output_ans = dict_info["TF_full"]
                        if len(input_question) <= 2048 and len(output_ans) <= 2048:
                            csvwriter.writerow([input_question, output_ans])
            print(f'Data has been written to {save_path}')
            len_csv = cr.count_non_empty_rows(save_path)
            print(f"len: {len_csv}")